# -*- coding:utf-8  -*-
# @Author: ChenYangMing
# @Time: 2024-08-30

import xgboost as xgb

from Common.config_utils import Config


class XGBoostModel:
    """
    XGBoost模型，用于分类或回归任务。
    :param objective: str，学习任务和相应的学习目标，例如 'binary:logistic'。
    :param num_estimators: int，要进行的提升轮数。
    :param learning_rate: float，学习率。
    :param max_depth: int，构建树时节点的最大深度。
    :param subsample: float，训练样本的抽样比。
    :param colsample_bytree: float，构建树时特征的采样比。
    :param gamma: float，定义节点分裂所需的最小损失函数下降值。
    """
    def __init__(
            self,
            objective='binary:logistic',
            num_estimators=Config.N_ESTIMATORS,
            learning_rate=Config.LR,
            max_depth=Config.MAX_DEPTH,
            subsample=Config.SUBSAMPLE,
            colsample_bytree=0.8,
            gamma=0,
    ):
        self.params = {
            'objective': objective,
            'num_estimators': num_estimators,
            'learning_rate': learning_rate,
            'max_depth': max_depth,
            'gamma': gamma,
            'subsample': subsample,
            'colsample_bytree': colsample_bytree
        }
        self.model = None

    def train(
            self,
            x_train,
            y_train,
            early_stopping_rounds=None,
    ):
        """
        训练 XGBoost 模型。
        :param x_train: numpy.ndarray 或 pandas.DataFrame，训练数据的特征。
        :param y_train: numpy.ndarray，训练数据的目标标签。
        :param early_stopping_rounds: int 或 None，早停法使用的轮数，如果为 None，则不使用早停。
        """
        xgb_matrix = xgb.DMatrix(x_train, label=y_train)
        if early_stopping_rounds:
            self.model = xgb.train(self.params, xgb_matrix, num_boost_round=self.params['num_estimators'],
                                   early_stopping_rounds=early_stopping_rounds, evals=[(xgb_matrix, 'train')])
        else:
            self.model = xgb.train(self.params, xgb_matrix, num_boost_round=self.params['num_estimators'])

    def predict(
            self,
            x_test,
    ):
        """
        使用训练好的模型进行预测。
        :param x_test: numpy.ndarray 或 pandas.DataFrame，测试数据的特征。
        :return: numpy.ndarray，模型预测的结果。
        """
        return self.model.predict(xgb.DMatrix(x_test))


# 使用示例
# xgb_model = XGBoostModel(objective='multi:softmax', num_estimators=100, learning_rate=0.1)  # 实例化 XGBoost 模型
# xgb_model.train(X_train, y_train, early_stopping_rounds=10)  # 训练模型
